import pickle
import torch
from functools import wraps
from optparse import OptionParser
import logging
import numpy as np


def process_data():
    BASE_DIR = './data/'

    word2ix = pickle.load(open(BASE_DIR + 'word_index.pkl', 'rb'))
    train_data = pickle.load(open(BASE_DIR + 'test_data.pkl', 'rb'))

    train_x = train_data[0]

    train_p = []
    i = 0
    for sent, y in zip(train_x, train_data[1]):
        sent.reverse()
        train_p.append((sent, y))
        i += 1
        print(i)

    pickle.dump(train_p, open(BASE_DIR + 'test_data_p.pkl', 'wb'))

    print('complete')


def partitions(number, k):
    '''
    Distribution of the folds
    Args:
        number: number of patients
        k: folds number
    '''
    n_partitions = np.ones(k) * int(number/k)
    n_partitions[0:(number % k)] += 1
    return n_partitions


def get_indices(n_splits = 3, subjects = 145, frames = 20):
    '''
    Indices of the set test
    Args:
        n_splits: folds number
        subjects: number of patients
        frames: length of the sequence of each patient
    '''
    l = partitions(subjects, n_splits)
    fold_sizes = l * frames
    indices = np.arange(subjects * frames).astype(int)
    current = 0
    for fold_size in fold_sizes:
        start = current
        stop =  current + fold_size
        current = stop
        yield(indices[int(start):int(stop)])

def k_folds(n_splits = 3, subjects = 145, frames = 20):
    '''
    Generates folds for cross validation
    Args:
        n_splits: folds number
        subjects: number of patients
        frames: length of the sequence of each patient
    '''
    indices = np.arange(subjects * frames).astype(int)
    for test_idx in get_indices(n_splits, subjects, frames):
        train_idx = np.setdiff1d(indices, test_idx)
    yield train_idx, test_idx


def get_user_paras():
        opt = OptionParser()
        opt.add_option('-b',
                       dest='batch',
                       type=int,
                       default=64,
                       help='batch_size')
        opt.add_option('-a',
                       dest="alpha",
                       type=float,
                       default=10,
                       help="bias alpha")
        opt.add_option('-g',
                       dest="gpu",
                       type=int,
                       default=0,
                       help="gpu no, default 0")
        opt.add_option('-e',
                       dest="epoch",
                       type=int,
                       default=100,
                       help="itre")
        opt.add_option('-r',
                       dest="regularization",
                       type=float,
                       default=0,
                       help="regularization, defaul 0")
        opt.add_option('--lr',
                       dest="lr",
                       type=float,
                       default=0.001,
                       help="learning rate, defaul 0.001")
        opt.add_option('--dropout',
                       dest="dropout",
                       type=float,
                       default=0.5,
                       help="add dropout to embedding layer")
        opt.add_option('-t',
                       dest="teacher",
                       type=int,
                       default=0,
                       help="use teacher forcing, default 1(True)")
        (options, args) = opt.parse_args()
        is_valid_paras = True
        error_messages = []
        batch = options.batch
        alpha = options.alpha
        gpu = options.gpu
        epoch = options.epoch
        regularization = options.regularization
        lr = options.lr
        dropout = options.dropout
        teacher = True if options.teacher else False

        if is_valid_paras:
            user_paras = {"batch": batch, "alpha": alpha, "gpu":gpu,
                          "epoch":epoch, "regularization":regularization,
                          "lr": lr, "dropout": dropout, "teacher": teacher}
            return user_paras

def get_res(pre_m, target, ix2tag):
    res = []
    for p_tag, t_tag in zip(pre_m.tolist(), target.tolist()):
        p_sent = []
        t_sent = []
        for p, t in zip(p_tag, t_tag):
            if p > 0:
                p_sent.append(ix2tag[p])
                t_sent.append(ix2tag[t])
        res.append([p_sent, t_sent])
    return res


def get_train_acc(pre_m, target):
    target = target.long()
    rel_tag_pre = torch.sum(pre_m.gt(0)).float()
    rel_tag_target = torch.sum(target.gt(0)).float()
    right_pre = (pre_m.eq(target) - pre_m.eq(0)).float().sum()

    if right_pre<0:
        right_pre = 0

    P = right_pre/rel_tag_pre if rel_tag_pre != 0 else 0
    R = right_pre/rel_tag_target
    F = (2 * P * R) / float(P + R) if P+R != 0 else 0

    result = torch.tensor([P, R, F])
    # print("pre_all: {}, right: {}, tar: {}".format(rel_tag_pre, right_pre, rel_tag_target))
    return result.data.tolist()


def get_train_tag_acc(pre_m, target):
    target = target.long()
    pre_m = pre_m.long()
    rel_tag_pre = torch.sum(pre_m.gt(1)).float()
    rel_tag_target = torch.sum(target.gt(1)).float()
    pre_excluded_one = pre_m.eq(target).float() - pre_m.eq(0).float() - pre_m.eq(1).float()
    right_pre = (pre_m.eq(target).float() - pre_m.eq(0).float() - pre_m.eq(1).float()).gt(0).sum().float()

    if right_pre<0:
        right_pre = 0

    P = right_pre/rel_tag_pre if rel_tag_pre != 0 else 0
    R = right_pre/rel_tag_target
    F = (2 * P * R) / float(P + R) if P+R != 0 else 0

    result = torch.tensor([P, R, F])
    # print("pre_all: {}, right: {}, tar: {}".format(rel_tag_pre, right_pre, rel_tag_target))
    return result.data.tolist()


def logged(func):
    log = logging.getLogger(func.__module__)
    @wraps(func)
    def wrapper(*args, **kwargs):
        log.log()
        result = func(*args, **kwargs)
        return result
    return wrapper


def find_pair(target, ix2dist):
    dist2ix = {v:k for k,v in ix2dist.items()}
    dist_gt_0 = []
    for sent in target:
        flag = [1 for i in range(len(sent))]
        tmp = []
        for i, t in enumerate(sent):
            if t>0 and flag[i] != 0:
                dist = ix2dist[t]
                pair_ix = sent.index(dist2ix[-dist])
                tmp.append((sent.index(t), pair_ix))
                flag[i] = 0
                flag[pair_ix] = 0
        dist_gt_0.append(tmp)

    return dist_gt_0


if __name__ == '__main__':
    process_data()